from ultralytics import YOLO
import argparse
from logger import Logger

log = Logger()

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    # cambricon dev id
    parser.add_argument("--device_id","--device_id", type=int, default=0, help="device_id")
    # MagicMind Model File
    parser.add_argument("--magicmind_model","--magicmind_model",type=str,default=None)
    # yolo task mode 
    parser.add_argument("--task", "--task", type=str, default="detect", help="'detect', 'segment', 'classify', 'pose')")
    # infer mode
    parser.add_argument("--val_or_predict", "--val_or_predict", type=str, default="val", help="'(val', 'predict')")
    # yolo model_cfg
    parser.add_argument("--model_cfg", "--model_cfg", type=str, default=None, help="model_cfg")
    # yolo data_cfg
    parser.add_argument("--data_cfg", "--data_cfg", type=str, default=None, help="data_cfg")
    # source
    parser.add_argument("--source", "--source", type=str, default=None, help="source")
    # model
    parser.add_argument("--model", "--model", type=str, default=None, help="model")
    # batch
    parser.add_argument("--batch", "--batch", type=int, default=16, help="batcj")

    args = parser.parse_args()

    model = YOLO(model = args.model_cfg,task = args.task)

    if args.val_or_predict == "val":
        model.val(
            data = args.data_cfg, 
            mm_model = args.magicmind_model, 
            mm_dev_id = args.device_id,
            batch = args.batch
            )

    elif args.val_or_predict == "predict":
        model.predict(
            source = args.source,
            model = args.model, 
            mm_model = args.magicmind_model,
            mm_dev_id = args.device_id,
            batch = args.batch
            )
    else:
        log.info("Error,Unsupport args.val_or_predict!")
